{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "gpt-2.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.8"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kc5cIgeEmv8o",
        "colab_type": "text"
      },
      "source": [
        "# Exporting GPT-2\n",
        "\n",
        "In this notebook, we'll show how to export [OpenAI's GPT-2 text generation model](https://github.com/openai/gpt-2) for serving."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RAWs29lAktOK",
        "colab_type": "text"
      },
      "source": [
        "First, we'll download the GPT-2 code repository:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gHs3aaFaLUXq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone --no-checkout https://github.com/openai/gpt-2.git\n",
        "!cd gpt-2 && git reset --hard ac5d52295f8a1c3856ea24fb239087cc1a3d1131"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A4al4P14nmni",
        "colab_type": "text"
      },
      "source": [
        "Next we'll specify the model size (choose one of 124M, 355M, or 774M):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3Y4bt6hkfuxY",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "import sys\n",
        "\n",
        "MODEL_SIZE = \"124M\" #@param {type:\"string\"}\n",
        "\n",
        "if MODEL_SIZE not in {\"124M\", \"355M\", \"774M\"}:\n",
        "    print(\"\\033[91m{}\\033[00m\".format('ERROR: MODEL_SIZE must be \"124M\", \"355M\", or \"774M\"'), file=sys.stderr)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C6xRx0Monh_j",
        "colab_type": "text"
      },
      "source": [
        "We can use `download_model.py` to download the model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kb50Z6NjbJBN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python3 ./gpt-2/download_model.py $MODEL_SIZE"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zz2ioOcpoPjV",
        "colab_type": "text"
      },
      "source": [
        "Next, we'll install the required packages:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vk4Q2RR-UZQm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install tensorflow==1.14.* numpy==1.* boto3==1.*"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KkVf5FmuUMrl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import sys\n",
        "import os\n",
        "import time\n",
        "import json\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Ay7qiQFoWRn",
        "colab_type": "text"
      },
      "source": [
        "Now we can export the model for serving:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GdnYXr1IKaF0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "sys.path.append(os.path.join(os.getcwd(), 'gpt-2/src'))\n",
        "import model, sample\n",
        "\n",
        "def export_for_serving(\n",
        "    model_name='124M',\n",
        "    seed=None,\n",
        "    batch_size=1,\n",
        "    length=None,\n",
        "    temperature=1,\n",
        "    top_k=0,\n",
        "    models_dir='models'\n",
        "):\n",
        "    \"\"\"\n",
        "    Export the model for TF Serving\n",
        "    :model_name=124M : String, which model to use\n",
        "    :seed=None : Integer seed for random number generators, fix seed to reproduce\n",
        "     results\n",
        "    :length=None : Number of tokens in generated text, if None (default), is\n",
        "     determined by model hyperparameters\n",
        "    :temperature=1 : Float value controlling randomness in boltzmann\n",
        "     distribution. Lower temperature results in less random completions. As the\n",
        "     temperature approaches zero, the model will become deterministic and\n",
        "     repetitive. Higher temperature results in more random completions.\n",
        "    :top_k=0 : Integer value controlling diversity. 1 means only 1 word is\n",
        "     considered for each step (token), resulting in deterministic completions,\n",
        "     while 40 means 40 words are considered at each step. 0 (default) is a\n",
        "     special setting meaning no restrictions. 40 generally is a good value.\n",
        "     :models_dir : path to parent folder containing model subfolders\n",
        "     (i.e. contains the <model_name> folder)\n",
        "    \"\"\"\n",
        "    models_dir = os.path.expanduser(os.path.expandvars(models_dir))\n",
        "\n",
        "    hparams = model.default_hparams()\n",
        "    with open(os.path.join(models_dir, model_name, 'hparams.json')) as f:\n",
        "        hparams.override_from_dict(json.load(f))\n",
        "\n",
        "    if length is None:\n",
        "        length = hparams.n_ctx\n",
        "    elif length > hparams.n_ctx:\n",
        "        raise ValueError(\"Can't get samples longer than window size: %s\" % hparams.n_ctx)\n",
        "\n",
        "    with tf.Session(graph=tf.Graph()) as sess:\n",
        "        context = tf.placeholder(tf.int32, [batch_size, None])\n",
        "        np.random.seed(seed)\n",
        "        tf.set_random_seed(seed)\n",
        "\n",
        "        output = sample.sample_sequence(\n",
        "            hparams=hparams, length=length,\n",
        "            context=context,\n",
        "            batch_size=batch_size,\n",
        "            temperature=temperature, top_k=top_k\n",
        "        )\n",
        "\n",
        "        saver = tf.train.Saver()\n",
        "        ckpt = tf.train.latest_checkpoint(os.path.join(models_dir, model_name))\n",
        "        saver.restore(sess, ckpt)\n",
        "\n",
        "        export_dir=os.path.join(models_dir, model_name, \"export\", str(time.time()).split('.')[0])\n",
        "        if not os.path.isdir(export_dir):\n",
        "            os.makedirs(export_dir)\n",
        "\n",
        "        builder = tf.saved_model.builder.SavedModelBuilder(export_dir)\n",
        "        signature = predict_signature_def(inputs={'context': context},\n",
        "        outputs={'sample': output})\n",
        "\n",
        "        builder.add_meta_graph_and_variables(sess,\n",
        "                                     [tf.saved_model.SERVING],\n",
        "                                     signature_def_map={\"predict\": signature},\n",
        "                                     strip_default_attrs=True)\n",
        "        builder.save()\n",
        "\n",
        "\n",
        "export_for_serving(top_k=40, length=256, model_name=MODEL_SIZE)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hGfSohMrowmg",
        "colab_type": "text"
      },
      "source": [
        "## Upload the model to AWS\n",
        "\n",
        "Cortex loads models from AWS, so we need to upload the exported model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BfB5QZ82ozj9",
        "colab_type": "text"
      },
      "source": [
        "Set these variables to configure your AWS credentials and model upload path:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B2RNuNk7o1c5",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "AWS_ACCESS_KEY_ID = \"\" #@param {type:\"string\"}\n",
        "AWS_SECRET_ACCESS_KEY = \"\" #@param {type:\"string\"}\n",
        "S3_UPLOAD_PATH = \"s3://my-bucket/text-generator/gpt-2\" #@param {type:\"string\"}\n",
        "\n",
        "import sys\n",
        "import re\n",
        "\n",
        "if AWS_ACCESS_KEY_ID == \"\":\n",
        "    print(\"\\033[91m {}\\033[00m\".format(\"ERROR: Please set AWS_ACCESS_KEY_ID\"), file=sys.stderr)\n",
        "\n",
        "elif AWS_SECRET_ACCESS_KEY == \"\":\n",
        "    print(\"\\033[91m {}\\033[00m\".format(\"ERROR: Please set AWS_SECRET_ACCESS_KEY\"), file=sys.stderr)\n",
        "\n",
        "else:\n",
        "    try:\n",
        "        bucket, key = re.match(\"s3://(.+?)/(.+)\", S3_UPLOAD_PATH).groups()\n",
        "    except:\n",
        "        print(\"\\033[91m {}\\033[00m\".format(\"ERROR: Invalid s3 path (should be of the form s3://my-bucket/path/to/file)\"), file=sys.stderr)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ics0omsrpS8V",
        "colab_type": "text"
      },
      "source": [
        "Upload the model to S3:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BnKncToppUhN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import boto3\n",
        "\n",
        "s3 = boto3.client(\"s3\", aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY)\n",
        "\n",
        "for dirpath, _, filenames in os.walk(\"models/{}/export\".format(MODEL_SIZE)):\n",
        "    for filename in filenames:\n",
        "        filepath = os.path.join(dirpath, filename)\n",
        "        filekey = os.path.join(key, MODEL_SIZE, filepath[len(\"models/{}/export/\".format(MODEL_SIZE)):])\n",
        "        print(\"Uploading s3://{}/{} ...\".format(bucket, filekey), end = '')\n",
        "        s3.upload_file(filepath, bucket, filekey)\n",
        "        print(\" ✓\")\n",
        "\n",
        "print(\"\\nUploaded model export directory to {}/{}\".format(S3_UPLOAD_PATH, MODEL_SIZE))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IIMVPhe2qkU4",
        "colab_type": "text"
      },
      "source": [
        "We also need to upload `vocab.bpe` and `encoder.json`, so that the encoder can encode the input text before making a request to the model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YdN8MtZxsO9V",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "print(\"Uploading s3://{}/{}/vocab.bpe ...\".format(bucket, key), end = '')\n",
        "s3.upload_file(os.path.join(\"models\", MODEL_SIZE, \"vocab.bpe\"), bucket, os.path.join(key, \"vocab.bpe\"))\n",
        "print(\" ✓\")\n",
        "\n",
        "print(\"Uploading s3://{}/{}/encoder.json ...\".format(bucket, key), end = '')\n",
        "s3.upload_file(os.path.join(\"models\", MODEL_SIZE, \"encoder.json\"), bucket, os.path.join(key, \"encoder.json\"))\n",
        "print(\" ✓\")"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}
